#include "myaes.h"
#include "string.h"
#include "stdlib.h"
#include "mbedtls/aes.h"

/**
  * @brief  ??PKCS#7??
  * @param  data: ???????
  * @param  len: ????
  * @param  block_size: ??? (AES?16)
  * @retval ???????
  */
size_t add_pkcs7_padding(unsigned char *data, size_t len, size_t block_size) {
    size_t padding_len = block_size - (len % block_size);
    size_t new_len = len + padding_len;
    
    // ??????
    for(size_t i = 0; i < padding_len; i++) {
        data[len + i] = (unsigned char)padding_len;
    }
    
    return new_len;
}

/**
  * @brief  ??PKCS#7??
  * @param  data: ??????
  * @param  len: ???? (?????????)
  * @retval ????????????
  */
size_t remove_pkcs7_padding(unsigned char *data, size_t len) {
    if(len == 0) return 0;
    
    // ???????
    unsigned char padding_len = data[len - 1];
    
    // ???????
    if(padding_len == 0 || padding_len > 16) {
        return len; // ????
    }
    
    // ????????????
    for(size_t i = len - padding_len; i < len; i++) {
        if(data[i] != padding_len) {
            return len; // ????
        }
    }
    
    return len - padding_len;
}

/**
  * @brief  AES-128-ECB ??(???????)
  * @param  key: 16????
  * @param  input: ????
  * @param  input_len: ??????
  * @param  output: ?????(?????)
  * @param  output_len: ???????
  * @retval 0??,??????
  */
int aes_128_ecb_encrypt(const unsigned char *key,
                        const unsigned char *input, size_t input_len,
                        unsigned char *output, size_t *output_len) {
    int ret;
    mbedtls_aes_context aes_ctx;
    
    // ???????
    size_t padded_len = input_len + (16 - (input_len % 16));
    unsigned char *padded_data = malloc(padded_len);
    if(!padded_data) {
        //return MBEDTLS_ERR_AES_ALLOC_FAILED;
				return -1;
    }
    
    // ?????????
    memcpy(padded_data, input, input_len);
    add_pkcs7_padding(padded_data, input_len, 16);
    
    // ???AES???
    mbedtls_aes_init(&aes_ctx);
    
    // ??????
    if((ret = mbedtls_aes_setkey_enc(&aes_ctx, key, 128)) != 0) {
        goto cleanup;
    }
    
    // ????
    for(size_t i = 0; i < padded_len; i += 16) {
        if((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT, 
                                       padded_data + i, output + i)) != 0) {
            goto cleanup;
        }
    }
    
    *output_len = padded_len;
    
cleanup:
    mbedtls_aes_free(&aes_ctx);
    free(padded_data);
    return ret;
}
												
/**
  * @brief  AES-128-ECB ??(???????)
  * @param  key: 16????
  * @param  input: ????
  * @param  input_len: ??????(???16???)
  * @param  output: ?????
  * @param  output_len: ???????(????)
  * @retval 0??,??????
  */
int aes_128_ecb_decrypt(const unsigned char *key,
                        const unsigned char *input, size_t input_len,
                        unsigned char *output, size_t *output_len) {
    int ret;
    mbedtls_aes_context aes_ctx;
    
    // ??????
    if(input_len % 16 != 0) {
        return MBEDTLS_ERR_AES_INVALID_INPUT_LENGTH;
    }
    
    // ???AES???
    mbedtls_aes_init(&aes_ctx);
    
    // ??????
    if((ret = mbedtls_aes_setkey_dec(&aes_ctx, key, 128)) != 0) {
        goto cleanup;
    }
    
    // ????
    for(size_t i = 0; i < input_len; i += 16) {
        if((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_DECRYPT, 
                                       input + i, output + i)) != 0) {
            goto cleanup;
        }
    }
    
    // ????
    *output_len = remove_pkcs7_padding(output, input_len);
    
cleanup:
    mbedtls_aes_free(&aes_ctx);
    return ret;
}
												
/**
  * @brief  AES-128-ECB ??(???????)
  * @param  key: 16????
  * @param  input: ????
  * @param  input_len: ??????
  * @param  output: ?????(?????)
  * @param  output_len: ???????
  * @retval 0??,??????
  */
int aes_128_ecb_encrypt_static(const unsigned char *key,
                               const unsigned char *input, size_t input_len,
                               unsigned char *output, size_t *output_len) {
    int ret;
    mbedtls_aes_context aes_ctx;
    
    // ???????
    size_t padded_len = input_len + (16 - (input_len % 16));
    
    // ??????????
    //if(padded_len > MAX_BUFFER_SIZE) {
		if(padded_len > 15) {
        //return MBEDTLS_ERR_AES_BUFFER_TOO_SMALL;
			  return -1;
    }
    
    // ??????????
    memcpy(output, input, input_len);
    
    // ????
    add_pkcs7_padding(output, input_len, 16);
    
    // ???AES???
    mbedtls_aes_init(&aes_ctx);
    
    // ??????
    if((ret = mbedtls_aes_setkey_enc(&aes_ctx, key, 128)) != 0) {
        goto cleanup;
    }
    
    // ????
    for(size_t i = 0; i < padded_len; i += 16) {
        if((ret = mbedtls_aes_crypt_ecb(&aes_ctx, MBEDTLS_AES_ENCRYPT, 
                                       output + i, output + i)) != 0) {
            goto cleanup;
        }
    }
    
    *output_len = padded_len;
    
cleanup:
    mbedtls_aes_free(&aes_ctx);
    return ret;
}
